package com.jxpanda.infrastructure.crypto.key.contract;

import com.jxpanda.infrastructure.crypto.toolkit.Base64Kit;

import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.security.*;
import java.security.spec.AlgorithmParameterSpec;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;

/**
 * 非对称密钥提供器：生成密钥对，加载公钥/私钥（仅支持 PEM/DER 格式）
 */
public interface KeyPairProvider<P extends AlgorithmParameterSpec> extends KeyProvider<KeyPair,P> {

    /**
     * 生成密钥对(使用默认配置生产，方便使用)
     */
    KeyPair generateKey();

    /**
     * 生成密钥对
     *
     * @param spec RSAKeyGenParameterSpec 或 ECGenParameterSpec
     */
    @Override
    default KeyPair generateKey(P spec) {
        try {
            KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance(getAlgorithm());
            keyPairGenerator.initialize(spec);
            return keyPairGenerator.generateKeyPair();
        } catch (NoSuchAlgorithmException | InvalidAlgorithmParameterException e) {
            throw new RuntimeException(e);
        }

    }

    default PublicKey loadPublicKey(byte[] bytes) {
        try {
            X509EncodedKeySpec spec = new X509EncodedKeySpec(bytes);
            KeyFactory keyFactory = KeyFactory.getInstance(getAlgorithm());
            return keyFactory.generatePublic(spec);
        } catch (NoSuchAlgorithmException | InvalidKeySpecException e) {
            throw new RuntimeException(e);
        }

    }

    /**
     * 从 Base64 或 PEM 格式字符串加载公钥
     */
    default PublicKey loadPublicKey(String publicKeyString) {
        byte[] bytes = Base64Kit.decode(stripPemHeaders(publicKeyString));
        return loadPublicKey(bytes);
    }

    /**
     * 从文件加载公钥（文件内容为 PEM/DER 的 Base64 编码）
     */
    default PublicKey loadPublicKeyFromFile(String path) {
        try {
            byte[] bytes = Files.readAllBytes(Paths.get(path));
            return loadPublicKey(bytes);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }


    default PrivateKey loadPrivateKey(byte[] bytes) {
        try {
            PKCS8EncodedKeySpec spec = new PKCS8EncodedKeySpec(bytes);
            KeyFactory keyFactory = KeyFactory.getInstance(getAlgorithm());
            return keyFactory.generatePrivate(spec);
        } catch (NoSuchAlgorithmException | InvalidKeySpecException e) {
            throw new RuntimeException(e);
        }

    }

    /**
     * 从 Base64 或 PEM 格式字符串加载私钥
     */
    default PrivateKey loadPrivateKey(String privateKeyString) {
        byte[] bytes = Base64Kit.decode(stripPemHeaders(privateKeyString));
        return loadPrivateKey(bytes);
    }

    /**
     * 从文件加载私钥（文件内容为 PEM/DER 的 Base64 编码）
     */
    default PrivateKey loadPrivateKeyFromFile(String path) {
        try {
            byte[] bytes = Files.readAllBytes(Paths.get(path));
            return loadPrivateKey(bytes);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 从字符串加载一对密钥
     */
    default KeyPair loadKeyPair(String publicKeyString, String privateKeyString) {
        PublicKey publicKey = loadPublicKey(publicKeyString);
        PrivateKey privateKey = loadPrivateKey(privateKeyString);
        return new KeyPair(publicKey, privateKey);
    }

    /**
     * 从文件加载一对密钥
     */
    default KeyPair loadKeyPairFromFiles(String publicKeyPath, String privateKeyPath) {
        PublicKey publicKey = loadPublicKeyFromFile(publicKeyPath);
        PrivateKey privateKey = loadPrivateKeyFromFile(privateKeyPath);
        return new KeyPair(publicKey, privateKey);
    }

    /**
     * 去除 PEM 格式头尾并清除空白
     */
    static String stripPemHeaders(String keyString) {
        return keyString
                .replaceAll("-----BEGIN [^-]+-----", "")
                .replaceAll("-----END [^-]+-----", "")
                .replaceAll("\\s+", "");
    }
}